import sys
import json
from datetime import datetime, timedelta
from PyQt5.QtWidgets import (QApplication, QMainWindow, QTableWidget, QTableWidgetItem,
                             QPushButton, QVBoxLayout, QWidget, QFileDialog, QComboBox,
                             QMessageBox)
from PyQt5.QtGui import QColor, QPixmap
from PyQt5.QtCore import Qt
import openpyxl


class ScheduleApp(QMainWindow):
    WEEKDAY_MAP = {
        "Monday": "周一",
        "Tuesday": "周二",
        "Wednesday": "周三",
        "Thursday": "周四",
        "Friday": "周五",
        "Saturday": "周六",
        "Sunday": "周日"
    }

    def __init__(self):
        super().__init__()
        self.shift_config = self.load_shift_config()
        self.initUI()

    def initUI(self):
        self.setWindowTitle("智能排班管理系统")
        self.setGeometry(100, 100, 1200, 600)

        # 主布局
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        layout = QVBoxLayout(main_widget)

        # 表格初始化
        self.table = QTableWidget()
        layout.addWidget(self.table)

        # 初始化表格结构
        self.set_table_headers()
        self.table.horizontalHeader().setDefaultSectionSize(120)
        self.table.verticalHeader().setDefaultSectionSize(40)
        self.table.setEditTriggers(QTableWidget.NoEditTriggers)

        # 操作按钮区
        btn_layout = QVBoxLayout()
        layout.addLayout(btn_layout)

        # 功能按钮
        self.reload_btn = QPushButton("重新加载数据 (F5)")
        self.reload_btn.setShortcut("F5")
        self.reload_btn.clicked.connect(self.reload_data)
        btn_layout.addWidget(self.reload_btn)

        self.export_btn = QPushButton("导出排班表 (Ctrl+E)")
        self.export_btn.setShortcut("Ctrl+E")
        self.export_btn.clicked.connect(self.export_to_excel)
        btn_layout.addWidget(self.export_btn)

        self.screenshot_btn = QPushButton("保存截图 (Ctrl+S)")
        self.screenshot_btn.setShortcut("Ctrl+S")
        self.screenshot_btn.clicked.connect(self.save_as_image)
        btn_layout.addWidget(self.screenshot_btn)

        # 初始化数据
        self.load_initial_data()

    def load_shift_config(self):
        """从配置文件加载班次设置"""
        try:
            with open("shift_config.json", "r", encoding="utf-8") as f:
                config = json.load(f)
                # 验证配置格式
                if "shifts" not in config or "colors" not in config:
                    raise ValueError("配置文件格式错误")
                return config
        except Exception as e:
            QMessageBox.critical(self, "配置错误",
                                 f"加载班次配置失败，使用默认配置\n错误信息：{str(e)}")
            return {
                "shifts": ["休息", "早班", "中班", "晚班", "请假"],
                "colors": {
                    "休息": "#F0F0F0",
                    "早班": "#FFD700",
                    "中班": "#98FB98",
                    "晚班": "#87CEEB",
                    "请假": "#FFA07A"
                }
            }

    def load_employee_data(self):
        """从Excel加载员工数据"""
        try:
            wb = openpyxl.load_workbook(r"f:\员工信息.xlsx")
            ws = wb.active
            return [cell.value for cell in ws['A'] if cell.value][1:]  # 跳过标题
        except Exception as e:
            QMessageBox.critical(self, "数据错误",
                                 f"无法读取员工信息文件\n错误信息：{str(e)}")
            return ["员工1", "员工2", "员工3"]  # 默认示例数据

    def set_table_headers(self):
        """动态生成完整周表头"""
        today = datetime.now()
        start_date = today - timedelta(days=today.weekday())

        headers = ["姓名"]
        for i in range(7):
            current_date = start_date + timedelta(days=i)
            date_str = current_date.strftime("%m/%d")
            weekday = current_date.strftime("%A")
            headers.append(f"{date_str}\n{self.WEEKDAY_MAP[weekday]}")

        self.table.setColumnCount(len(headers))
        self.table.setHorizontalHeaderLabels(headers)

    def create_shift_combobox(self):
        """创建可配置的班次下拉框"""
        combo = QComboBox()
        combo.addItems(self.shift_config["shifts"])
        combo.setCurrentIndex(0)
        self.update_combo_style(combo, combo.currentText())
        combo.currentIndexChanged.connect(
            lambda: self.update_combo_style(combo, combo.currentText())
        )
        return combo

    def update_combo_style(self, combo, shift):
        """更新下拉框样式"""
        color = self.shift_config["colors"].get(shift, "#FFFFFF")
        combo.setStyleSheet(f"""
            QComboBox {{
                background-color: {color};
                border: 1px solid #D3D3D3;
                padding: 3px;
            }}
            QComboBox::drop-down {{
                width: 20px;
                border-left: 1px solid #D3D3D3;
            }}
        """)

    def load_initial_data(self):
        """初始化表格数据"""
        employees = self.load_employee_data()
        self.table.setRowCount(len(employees))

        for row, name in enumerate(employees):
            # 员工姓名
            self.table.setItem(row, 0, QTableWidgetItem(str(name)))

            # 初始化班次单元格
            for col in range(1, self.table.columnCount()):
                combo = self.create_shift_combobox()
                self.table.setCellWidget(row, col, combo)

    def reload_data(self):
        """重新加载所有数据"""
        self.shift_config = self.load_shift_config()
        self.load_initial_data()
        QMessageBox.information(self, "系统提示", "数据和配置已重新加载")

    def export_to_excel(self):
        """导出为Excel文件"""
        path, _ = QFileDialog.getSaveFileName(
            self, "保存文件", "", "Excel文件 (*.xlsx)")
        if not path:
            return

        wb = openpyxl.Workbook()
        ws = wb.active

        # 写入表头
        headers = [self.table.horizontalHeaderItem(i).text().replace("\n", " ")
                   for i in range(self.table.columnCount())]
        ws.append(headers)

        # 写入数据
        for row in range(self.table.rowCount()):
            row_data = [self.table.item(row, 0).text()]
            for col in range(1, self.table.columnCount()):
                combo = self.table.cellWidget(row, col)
                row_data.append(combo.currentText())
            ws.append(row_data)

        # 调整列宽
        for col in range(1, self.table.columnCount() + 1):
            ws.column_dimensions[openpyxl.utils.get_column_letter(col)].width = 15

        wb.save(path)

    def save_as_image(self):
        """保存表格截图"""
        path, _ = QFileDialog.getSaveFileName(
            self, "保存截图", "", "PNG图像 (*.png)")
        if not path:
            return

        # 计算内容尺寸
        header_height = self.table.horizontalHeader().height()
        row_height = sum([self.table.rowHeight(i) for i in range(self.table.rowCount())])
        total_width = self.table.horizontalHeader().length()
        total_height = header_height + row_height

        # 渲染截图
        pixmap = QPixmap(total_width, total_height)
        self.table.render(pixmap)
        pixmap.save(path)


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = ScheduleApp()
    window.show()
    sys.exit(app.exec_())
